package layers

type maxPull2d struct {
	poolSize uint
	stride   uint
	count    uint
	inpDim   uint
	outDim   uint
	memory   [][][]coord
}

type coord struct {
	x uint
	y uint
}

// 6 28 2 2
func NewMaxPull2d(inpCount, inpDim uint, poolSize uint, stride uint) *maxPull2d {
	return &maxPull2d{
		poolSize: 2, //not inmlemented
		stride:   stride,

		count:  inpCount,
		inpDim: inpDim,
		outDim: inpDim / poolSize,
		memory: nil,
	}
}

func (l *maxPull2d) Forward(x []float64) []float64 {
	//log.Println("max")
	//log.Println("l.count:", l.count, "l.intDim:", l.inpDim, "len(x):", len(x))
	tmp := ConvertTensor1to2(x, l.count) // получаем входные каналы
	var result [][]float64
	for i, v := range tmp {

		y := ConvertTensor1to2(v, l.inpDim) // получаем матрицу канала
		res := make([][]float64, 0)
		l.memory = append(l.memory, make([][]coord, 0))

		for iy, iRes := 0, 0; iy < int(l.inpDim); iy, iRes = iy+int(l.stride), iRes+1 { //
			res = append(res, make([]float64, 0))
			l.memory[i] = append(l.memory[i], make([]coord, 0))
			for jy, jRes := 0, 0; jy < int(l.inpDim); jy, jRes = jy+int(l.stride), jRes+1 {
				tiy, tjy := iy, jy
				val, coord := max(y[tiy][tjy], y[tiy][tjy+1], y[tiy+1][tjy], y[tiy+1][tjy+1], uint(tiy), uint(tjy))
				res[iRes] = append(res[iRes], val)
				tcoord := coord
				l.memory[i][iRes] = append(l.memory[i][iRes], tcoord)
			}
		}
		//log.Println("res len:", len(res), "x", len(res[0]))
		result = append(result, ConvertTensor2to1(res))
		//log.Println("result len", len(result[len(result)-1]))
	}
	//log.Println("size out:", len(result), "x", len(result[0]))
	//log.Println("end max")
	//log.Println(l.memory)
	return ConvertTensor2to1(result)
}
func (l *maxPull2d) BackProp(err []float64) []float64 {
	//log.Println("back max")
	convErr := ConvertTensor1to2(err, l.count)
	var result [][]float64
	for iv, v := range convErr {
		//log.Println(l.memory[iv])
		//errEl := ConvertTensor1to2(err, l.outDim)
		out := make([][]float64, l.inpDim)
		for i := 0; i < int(l.inpDim); i++ {
			out[i] = make([]float64, l.inpDim)
		}

		err2d := ConvertTensor1to2(v, l.outDim)
		for i := 0; i < int(l.outDim); i++ {
			for j := 0; j < int(l.outDim); j++ {
				pos := l.memory[iv][i][j]
				out[pos.x][pos.y] += err2d[i][j]
				//log.Println(pos.x, pos.y)
			}
		}
		//log.Println("!!!!!!!!!!!!\n", out)
		//log.Panicln(out)
		result = append(result, ConvertTensor2to1(out))
	}
	//log.Println("end back max")
	//log.Println(len(result), len(result[0]))
	return ConvertTensor2to1(result)
}

func max(a, b, c, d float64, startPosX, startPosY uint) (float64, coord) {
	xx, yy := startPosX, startPosY
	if a >= b && a >= c && a >= d {
		return a, coord{x: xx, y: yy}
	} else if b >= a && b >= c && b >= d {
		return b, coord{x: xx, y: yy + 1}
	} else if c >= a && c >= b && c >= d {
		return c, coord{x: xx + 1, y: yy}
	} else {
		return d, coord{x: xx + 1, y: yy + 1}
	}

}
